{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  HMinCaml.Alpha
-- Copyright   :  (c) Stephen Tetley 2012
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Alpha renaming / conversion.
--
--------------------------------------------------------------------------------

module HMinCaml.Alpha
  (

    alpha
  , alpha'

  ) where

import qualified HMinCaml.HMonad as H
import HMinCaml.Id
import HMinCaml.KNormal
import qualified HMinCaml.Type as T


import Control.Applicative
import Control.Monad
import qualified Data.Map as M



--------------------------------------------------------------------------------
-- alpha renaming

type Env = M.Map IdS IdS

newtype Mon a = Mon { getMon :: Env -> H.HMonad a }


instance Functor Mon where
  fmap f ma = Mon $ \r -> fmap f $ getMon ma r

instance Applicative Mon where
  pure a    = Mon $ \_ -> pure a
  af <*> av = Mon $ \r -> getMon af r <*> getMon av r

instance Monad Mon where
  return    = pure
  ma >>= k  = Mon $ \r -> getMon ma r >>= \a -> getMon (k a) r

instance MonadPlus Mon where
  mzero         = Mon $ \_ -> mzero
  ma `mplus` mb = Mon $ \r -> getMon ma r `mplus` getMon mb r

-- Env must be parametric
--
runMon :: Env -> Mon a -> H.HMonad a
runMon env ma = getMon ma env

local :: (Env -> Env) -> Mon a -> Mon a
local f ma = Mon $ \r -> getMon ma (f r)


genid :: IdS -> Mon IdS
genid ss = Mon $ \_ -> H.genid ss

localBinding :: IdS -> IdS -> Mon a -> Mon a
localBinding v v1 ma = local (M.insert v v1) ma

localBindings :: [(IdS,IdS)] -> Mon a -> Mon a
localBindings vts ma = Mon $ \r -> getMon ma (foldr fn r vts)
  where
    fn (v,t) m = M.insert v t m

find :: IdS -> Mon IdS
find x = Mon $ \r -> maybe (return x) return $ M.lookup x r
    
findArgs :: [(IdS,T.Type)] -> Mon [(IdS,T.Type)]
findArgs = mapM findarg 
  where
    findarg (w,ty) = (\a -> (a,ty)) <$> find w


alpha :: Expr -> H.HMonad Expr
alpha expr = runMon M.empty $ ac expr

alpha' :: Env -> Expr -> H.HMonad Expr
alpha' env expr = runMon env $ ac expr


ac :: Expr -> Mon Expr
ac (Unit)                 = return Unit
ac (Int i)                = return $ Int i
ac (Float d)              = return $ Float d
ac (Neg x)                = Neg  <$> find x
ac (Add x y)              = Add  <$> find x <*> find y
ac (Sub x y)              = Sub  <$> find x <*> find y
ac (FNeg x)               = FNeg <$> find x
ac (FAdd x y)             = FAdd <$> find x <*> find y
ac (FSub x y)             = FSub <$> find x <*> find y
ac (FMul x y)             = FMul <$> find x <*> find y
ac (FDiv x y)             = FDiv <$> find x <*> find y
ac (IfEq x y e1 e2)       = IfEq <$> find x <*> find y <*> ac e1 <*> ac e2
ac (IfLE x y e1 e2)       = IfLE <$> find x <*> find y <*> ac e1 <*> ac e2
ac (Let (x,t) e1 e2)      = do 
    x1 <- genid x
    ea <- ac e1
    eb <- localBinding x x1 $ ac e2
    return $ Let (x1,t) ea eb

ac (Var x)                = Var <$> find x

ac (LetRec (Fundef (x,t) yts e1) e2) = do
    x1 <- genid x
    localBinding x x1 $ do 
        let ys = map fst yts
        ys1 <- mapM genid ys
        fd  <- localBindings (zip ys ys1) $ 
                (\name args ea ->  Fundef (name,t) args ea)
                    <$> find x <*> findArgs yts <*> ac e1
        eb  <- ac e2
        return $ LetRec fd eb

ac (App x ys)               = App <$> find x <*> mapM find ys
ac (Tuple xs)               = Tuple <$> mapM find xs

ac (LetTuple xts y e)       = do
    let xs = map fst xts
    xs1 <- mapM genid xs
    y1 <- find y
    localBindings (zip xs xs1) $ 
       (\args ea -> LetTuple args y1 ea) <$> findArgs xts <*> ac e

ac (Get x y)                = Get <$> find x <*> find y
ac (Put x y z)              = Put <$> find x <*> find y <*> find z
ac (ExtArray x)             = ExtArray <$> find x
ac (ExtFunApp x ys)         = ExtFunApp <$> find x <*> mapM find ys


    